import os
from glob import glob
from tqdm import tqdm
from typing import Optional
from paddleocr import PaddleOCR

ocr = PaddleOCR(lang="ch", show_log=False, mode="PP-OCRv4")

# init dir

if not os.path.exists("./data"):
    os.mkdir("./data")
if not os.path.exists("./data/success"):
    os.mkdir("./data/success")
if not os.path.exists("./data/failed"):
    os.mkdir("./data/failed")
# img_path = "./crop_img/638806_63102748925.jpg"
# result = ocr.ocr(img_path, cls=True)
# print(result)
# for one in result[0]:
#     print(one)
"""
[[[79.0, 125.0], [158.0, 125.0], [158.0, 143.0], [79.0, 143.0]], ('限时到手价', 0.9994503259658813)]
[[[221.0, 125.0], [263.0, 125.0], [263.0, 148.0], [221.0, 148.0]], ('活动', 0.9986187219619751)]
[[[48.0, 150.0], [192.0, 156.0], [190.0, 201.0], [46.0, 195.0]], ('418.6', 0.9996075630187988)]

"""

a = [[[48.0, 150.0], [192.0, 156.0], [190.0, 201.0], [46.0, 195.0]], ('418.6', 0.9996075630187988)]


def ocr_price(img_path: str) -> Optional[float]:
    result = ocr.ocr(img_path, cls=True)
    if not result or not result[0]:
        return None
    for one in result[0]:
        ocr_tuple = one[1]
        ocr_string, confidence = ocr_tuple
        try:
            ocr_string = ocr_string.replace("￥", "").strip()
            ocr_string = ocr_string.replace("¥", "").strip()
            ocr_string = ocr_string.replace("*", "").strip()
            if ocr_string.count(".") > 1:
                ocr_string = ocr_string.split(".")[0]

            ocr_float = float(ocr_string)
            return ocr_float
        except Exception as e:
            pass
    return None


ocr_field = 0


def to_new_jpg(idx: int, img_path: str, ocr_float: float = None):
    with open(img_path, "rb") as f:
        content = f.read()
    if ocr_float:
        with open(f"./data/success/{idx}_{ocr_float}.jpg", "wb") as f:
            f.write(content)
    else:
        with open(f"./data/failed/{idx}_{ocr_float}.jpg", "wb") as f:
            f.write(content)


img_glob_path = "./crop_img/*.jpg"
img_list = glob(img_glob_path)

all_idx = 0
for img_path in tqdm(img_list):
    all_idx += 1
    ocr_float = ocr_price(img_path)
    to_new_jpg(all_idx, img_path, ocr_float)

print("all done")